from numpy import empty, load
from math import sqrt,acos,sin,cos
import os,sys

from enthought.tvtk.api import tvtk,write_data
from numpy import *

def write_vtkxml_sgrid(data, filename):
    if type(filename) == type(''):
        file = open(filename,'w')
    for k in data.keys():
        data[k] = data[k].astype('f')
    x,y,z = [data[i] for i in ('x','y','z')]
    dims = x.shape
    var = [i for i in data.keys() if i not in ('x','y','z')]
    sg = tvtk.StructuredGrid(points=c_[z.flat,y.flat,x.flat],dimensions=array(dims))
    pd = tvtk.PointData()
    for v in var:
        pd.scalars = ravel(data[v].astype('f'))
        pd.scalars.name = v
        sg.point_data.add_array(pd.scalars)
    if 'u' in var and 'v' in var and 'w' in var:
        print 'Detected velocity vector'
        u,v,w = [data[i] for i in ('u','v','w')]
        sg.point_data.vectors = c_[ravel(w),ravel(v),ravel(u)]
        sg.point_data.vectors.name = 'V'
    xsgw = tvtk.XMLStructuredGridWriter(input=sg, file_name=filename, data_mode='binary')
    xsgw.write()

def save_img(t,nprocs,datafolder,resultsfolder):
    """
    Saves image obtained at time t. Obtains required data from datafolder and 
    saves image in resultsfolder.
    """
    # combine arrays from all procs
    x,y,z,gx,gy,gz,vx,vy,vz = [],[],[],[],[],[],[],[],[]
    for currentproc in range(nprocs):
        filename = str(currentproc)+'_t='+str(round(t,4))+'_sec'
        filename = os.path.join(datafolder,gridfile)
        fb = open(filename,'r')
        data = f.readlines()
        x += [float(number) for number in data[0].strip("'=','[',']','x','\n'").split(',')]
        y += [float(number) for number in data[1].strip("'=','[',']','y','\n'").split(',')]
        z += [float(number) for number in data[2].strip("'=','[',']','z','\n'").split(',')]
        gx += [float(number) for number in data[3].strip("'=','[',']','gx','\n'").split(',')]
        gy += [float(number) for number in data[4].strip("'=','[',']','gy','\n'").split(',')]
        gz += [float(number) for number in data[5].strip("'=','[',']','gz','\n'").split(',')]
        vx += [float(number) for number in data[6].strip("'=','[',']','vx','\n'").split(',')]
        vy += [float(number) for number in data[7].strip("'=','[',']','vy','\n'").split(',')]
        vz += [float(number) for number in data[8].strip("'=','[',']','vz','\n'").split(',')]
    
    # obtained arrays, now plot them
    #s = [sqrt(gx[i]**2+gy[i]**2+gz[i]**2) for i in range(len(gx))]
    print 'yo1'
    print len(x)
    #mlab.points3d(x,y,z,s,scale_mode='none',scale_factor=0.05,colormap='jet')
    mlab.points3d(x,y,z,scale_mode='none',scale_factor=0.05,colormap='jet')
    mlab.quiver3d(x,y,z,vx,vy,vz)
    #mlab.flow(x,y,z,vx,vy,vz)
    print 'yo'
    #mlab.view(90,90)
    #mlab.roll(90)
    #a = mlab.axes()
    #a.axes.use_data_bounds = False
    #a.axes.bounds = [-1.5, 1.5, -1.5, 1.5, 0, 0.5]
    #img_filename = os.path.join(resultsfolder,'t='+str(round(t,2))+'_sec.png')
    #mlab.savefig(img_filename)
    mlab.show()
    mlab.clf()
    
def save_img_bg(t,nprocs,datafolder,resultsfolder):
    """
    Saves image obtained at time t. Obtains required data from datafolder and 
    saves image in resultsfolder.
    """
    # combine arrays from all procs
    xb,yb,zb,gxb,gyb,gzb,vxb,vyb,vzb = [],[],[],[],[],[],[],[],[]
    xg,yg,zg,gxg,gyg,gzg,vxg,vyg,vzg = [],[],[],[],[],[],[],[],[]
    for currentproc in range(nprocs):
        blobfile = str(currentproc)+'b_t='+str(round(t,4))+'_sec'
        gridfile = str(currentproc)+'g_t='+str(round(t,4))+'_sec'
        blobfile = os.path.join(datafolder,blobfile)
        gridfile = os.path.join(datafolder,gridfile)
        
        # reading blob data
        fb = load(blobfile)
        xb,yb,zb,gxb,gyb,gzb,vxb,vyb,vzb = (fb['x'],fb['y'],fb['z'],
                                            fb['gx'],fb['gy'],fb['gz'],
                                            fb['vx'],fb['vy'],fb['vz'])
        # reading grid data
        fg = load(gridfile)
        xg,yg,zg,gxg,gyg,gzg,vxg,vyg,vzg = (fg['x'],fg['y'],fg['z'],
                                            fg['gx'],fg['gy'],fg['gz'],
                                            fg['vx'],fg['vy'],fg['vz'])
    
    # add grid data
    num = (25,25,25)
    xg3d = empty(num)
    yg3d = empty(num)
    zg3d = empty(num)
    vxg3d = empty(num)
    vyg3d = empty(num)
    vzg3d = empty(num)
    xg3d.flat[:] = xg
    yg3d.flat[:] = yg
    zg3d.flat[:] = zg
    vxg3d.flat[:] = vxg
    vyg3d.flat[:] = vyg
    vzg3d.flat[:] = vzg
    
    # save grid data
    data_g = {'x':xg3d,'y':yg3d,'z':zg3d,'u':vxg3d,'v':vyg3d,'w':vzg3d}
    write_vtkxml_sgrid(data_g,'vel_field.vts')
    savez('vel_field.npz', **data_g)
    
    # add point data
    pd = tvtk.PolyData(points=c_[xb,yb,zb])
    
    # add vor data to points
    vors = c_[gxb,gyb,gzb]
    pvors = tvtk.DoubleArray()
    pvors.name = 'vor'
    pvors.from_array(vors)
    pd.point_data.add_array(pvors)
    
    # add vel data to points
    vels = c_[vxb,vyb,vzb]
    pvels = tvtk.DoubleArray()
    pvels.name = 'vel'
    pvels.from_array(vels)
    pd.point_data.add_array(pvels)
    
    # save point data
    write_data(pd,'blobs.vtk')
    data_b = {'x':xb,'y':yb,'z':zb,'gx':gxb,'gy':gyb,'gz':gzb,'vx':vxb,'vy':vyb,'vz':vzb}
    savez('blobs.npz',**data_b)
    
    print 'saved vec data'
    
    from enthought.mayavi import mlab
    mlab.clf()
    
    # plot flow
    #mlab.flow(xg3d,yg3d,zg3d,vxg3d,vyg3d,vzg3d)
    
    # plot blobs
    sb = [sqrt(gxb[i]**2+gyb[i]**2+gzb[i]**2) for i in range(len(gxb))]
    print 'yo1'
    #mlab.points3d(xb,yb,zb,sb,scale_mode='none',scale_factor=0.05,colormap='jet')
    #mlab.points3d(x,y,z,scale_mode='none',scale_factor=0.05,colormap='jet')
    mlab.quiver3d(xb,yb,zb,vxb,vyb,vzb,extent=[-0.1,0.1,-0.1,0.1,-0.1,0.1])
    #mlab.flow(x,y,z,vx,vy,vz)
    print 'yo'
    #mlab.view(90,90)
    #mlab.roll(90)
    #a = mlab.axes()
    #a.axes.use_data_bounds = False
    #a.axes.bounds = [-1.5, 1.5, -1.5, 1.5, 0, 0.5]
    #img_filename = os.path.join(resultsfolder,'t='+str(round(t,2))+'_sec.png')
    #mlab.savefig(img_filename)
    #execfile('mayavi_script.py')
    mlab.show()

def save_img1(t,nprocs,datafolder,datafolder_d,resultsfolder):
    """
    This is exactly same as save_img except that it also plots particle 
    positions from direct calculation.
    """
    x,y,z,gx,gy,gz,c = [],[],[],[],[],[],[]
    for dfolder in [datafolder,datafolder_d]:
        for currentproc in range(nprocs):
            filename = str(currentproc)+'_t='+str(round(t,3))+'_sec'
            filename = os.path.join(dfolder,filename)
            f = open(filename,'r')
            data = f.readlines()
            x += [float(number) for number in data[0].strip("'=','[',']','x','\n'").split(',')]
            y += [float(number) for number in data[1].strip("'=','[',']','y','\n'").split(',')]
            z += [float(number) for number in data[2].strip("'=','[',']','z','\n'").split(',')]
            gx += [float(number) for number in data[3].strip("'=','[',']','gx','\n'").split(',')]
            gy += [float(number) for number in data[4].strip("'=','[',']','gy','\n'").split(',')]
            gz += [float(number) for number in data[5].strip("'=','[',']','gz','\n'").split(',')]
            color = 1 if dfolder is datafolder else 0 # direct: blue, tree: red
            c += [color for number in data[0].strip("'=','[',']','x','\n'").split(',')]

    # obtained arrays, now plot them
    s = [sqrt(gx[i]**2+gy[i]**2+gz[i]**2) for i in range(len(gx))]
    print len(x)
    mlab.points3d(x,y,z,c,scale_mode='none',scale_factor=0.05,colormap='jet')
    print 'yo'
    mlab.view(90,90)
    mlab.roll(90)
    a = mlab.axes()
    a.axes.use_data_bounds = False
    a.axes.bounds = [-1.5, 1.5, -1.5, 1.5, 0, 0.5]
    img_filename = os.path.join(resultsfolder,'t='+str(round(t,3))+'_sec.png')
    mlab.savefig(img_filename)
    mlab.clf()
    
def ideal_flow(t,nprocs,datafolder,resultsfolder):
    """
    Saves image obtained at time t. Obtains required data from datafolder and 
    saves image in resultsfolder.
    """
    currentproc = 0
    # combine arrays from all procs
    x,y,z,gx,gy,gz,vx,vy,vz = [],[],[],[],[],[],[],[],[]
    filename = str(currentproc)+'_t='+str(round(t,4))+'_sec'
    filename = os.path.join(datafolder,filename)
    f = open(filename,'r')
    data = f.readlines()
    x = [float(number) for number in data[0].strip("'=','[',']','x','\n'").split(',')]
    y = [float(number) for number in data[1].strip("'=','[',']','y','\n'").split(',')]
    z = [float(number) for number in data[2].strip("'=','[',']','z','\n'").split(',')]
    for i in range(len(x)):
        r = sqrt(x[i]**2 + y[i]**2 + z[i]**2)
        theta = acos(z[i]/r)
        phi = acos(x[i]/(sqrt(x[i]**2 + y[i]**2)))
        a = 0.01
        u0 = 7.3e-5
        ur = u0*cos(theta)*(1. - 1.5*a/r + 0.5*(a/r)**3)
        utheta = -u0*sin(theta)*(1. - 0.75*a/r - 0.25*(a/r)**3)
        uphi = 0.
        
        vx += [ur*sin(theta)*cos(phi) + utheta*cos(theta)*cos(phi)]
        vy += [ur*sin(theta)*sin(phi) + utheta*cos(theta)*sin(phi)]
        vz += [ur*cos(theta) - utheta*sin(theta)]

    print 'yo1'
    #mlab.points3d(x,y,z,s,scale_mode='none',scale_factor=0.05,colormap='jet')
    mlab.quiver3d(x,y,z,vx,vy,vz)
    print 'yo'
    #mlab.view(90,90)
    #mlab.roll(90)
    #a = mlab.axes()
    #a.axes.use_data_bounds = False
    #a.axes.bounds = [-1.5, 1.5, -1.5, 1.5, 0, 0.5]
    #img_filename = os.path.join(resultsfolder,'t='+str(round(t,2))+'_sec.png')
    #mlab.savefig(img_filename)
    mlab.show()
    mlab.clf()


if __name__=="__main__":
    nprocs = 1
    dt = 0.001
    tf = 0.01
    casename = sys.argv[1]
    casedir = os.path.join(os.path.abspath('../../examples'),casename)
    datafolder = os.path.join(casedir,'raw')
    resultsfolder = os.path.join(casedir,'results')
    t = 0
    while t<tf:
        #ideal_flow(t,nprocs,datafolder,resultsfolder)
        save_img_bg(t,nprocs,datafolder,resultsfolder)
        t += dt
